import datetime
import os
import numpy as np
import pandas as pd
import Core.Config as Config
import Core.Gadget as Gadget
from Utility.IndustryAnalysis import get_stock_with_industry
from Utility.industry_builder import IndustryBuilder


# 计算指数的市值权重
def get_index_constituent_weight(database, index_symbol, update_datetime):

    # 确认指数成分日期
    s_update_datetime = update_datetime.strftime('%Y-%m-%d')
    date_list = database.ExecuteSQL("financial_data", "SELECT distinct(date) FROM  financial_data.index_constituents where index_symbol='" + index_symbol + "' and date <= '" + s_update_datetime + "' order by date")
    max_date = date_list[-1][0]

    #
    filter = [("index_symbol", index_symbol), ("date", max_date)]
    df_constituent = database.GetDataFrame("financial_data", "index_constituents", filter=filter, projection=["symbol", "weight", "date"])

    return df_constituent


# 计算全市场-行业权重
def calc_market_industry_weight(database, profile_datetime, industry_type="sw_industry1", use_free_float_cap=False):
    # 读取市值信息
    df_market_stock_weight = calc_market_stock_weight(database, profile_datetime)

    # 读取行业信息
    df_stock_list = database.GetDataFrame("financial_data", "stock_instrument", projection=["symbol", "description", industry_type])
    df_stock_list.rename(columns={industry_type: "industry_name"}, inplace=True)

    # 补充行业信息
    df_market_stock_weight = pd.merge(df_market_stock_weight, df_stock_list[["symbol","description","industry_name"]], how="left", on="symbol")  # 补充股票行业信息

    # 补充行业代码
    industry_builder = IndustryBuilder(database=database, industry_type=industry_type)
    df_industry = industry_builder.get_industry_mapping()
    df_industry.rename(columns={industry_type + "_name": "industry_name", industry_type + "_symbol": "industry_symbol"}, inplace=True)

    #
    df_market_stock_weight = pd.merge(df_market_stock_weight, df_industry, how="left", on="industry_name")

    # 行业聚合
    df_market_industry_weight = df_market_stock_weight.groupby(["industry_name"]).agg({"industry_symbol": "first", 'stock_cap': np.sum})
    df_market_industry_weight.rename(columns={"stock_cap": "industry_cap"}, inplace=True)

    # 恢复index字段
    df_market_industry_weight.reset_index(inplace=True)

    # 计算权重
    df_market_industry_weight["market_weight"] = df_market_industry_weight["industry_cap"] / df_market_industry_weight["industry_cap"].sum()

    #
    return df_market_industry_weight


# 计算全市场-行业权重
# 修改部分底层函数，为了不破坏原有数据，同时保留新旧函数
def calc_market_industry_weight_v2(database, profile_datetime, industry_type="sw_industry1", use_free_float_cap=False):
    # 读取市值信息
    df_market_stock_weight = calc_market_stock_weight(database, profile_datetime)

    # 读取行业信息
    # 老方法
    # df_stock_list = database.GetDataFrame("financial_data", "stock_instrument", projection=["symbol", "description", "sw_industry1"])
    # df_stock_list.rename(columns={"sw_industry1": "industry"}, inplace=True)
    # 新方法
    df_stock_list = get_stock_with_industry(database, industry_type, profile_datetime)

    # 补充行业信息
    df_market_stock_weight = pd.merge(df_market_stock_weight, df_stock_list[["symbol","industry"]], how="left", on="symbol")  # 补充股票行业信息

    # 行业聚合
    df_market_industry_weight = df_market_stock_weight.groupby(["industry"]).agg({'stock_cap': np.sum})
    df_market_industry_weight.rename(columns={"stock_cap": "industry_cap"}, inplace=True)

    # 计算权重
    df_market_industry_weight["market_weight"] = df_market_industry_weight["industry_cap"] / df_market_industry_weight["industry_cap"].sum()

    #
    return df_market_industry_weight


# 计算全市场-股票对应权重
def calc_market_stock_weight(database, profile_datetime, use_free_float_cap=False):
    #
    s_update_datetime = profile_datetime.strftime('%Y-%m-%d')
    sql_str = "select max(date) from financial_data.stock_dailybar where date <=" + "'" + s_update_datetime + "'"
    max_date = database.ExecuteSQL("financial_data", sql_str)
    max_date = max_date[0][0]

    df_total_market = database.GetDataFrame("financial_data", "stock_dailybar", filter=[("date", max_date)], projection=["symbol", "close", "total_shares", "free_float_shares"])

    if use_free_float_cap:  # 自由流通市值加权
        df_total_market["stock_cap"] = df_total_market["close"] * df_total_market["free_float_shares"]

    else:  # 市值加权
        df_total_market["stock_cap"] = df_total_market["close"] * df_total_market["total_shares"]
    #
    total_stock_cap = df_total_market["stock_cap"].sum()
    df_total_market["market_weight"] = df_total_market["stock_cap"] / total_stock_cap
    #
    return df_total_market



#
if __name__ == '__main__':
    #
    path_filename = os.getcwd() + r"\..\Config\config_local.json"
    database = Config.create_database(database_type="MySQL", config_file=path_filename, config_field="MySQL")


    update_datetime = datetime.datetime(2022,8,1)
    index_symbol = "000905.SH"
    # get_index_constituent_weight(database, index_symbol, update_datetime)

    update_datetime = datetime.datetime(2022,10,31)
    df_market_industry_weight = calc_market_industry_weight(database, update_datetime, industry_type="citics_industry1", use_free_float_cap=False)
    for index , row in df_market_industry_weight.iterrows():
        print(index, row["industry_symbol"], row["industry_cap"], row["market_weight"])
